# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import logging
import sys

from tblib import pickling_support

from coin.util.slurm.base import (pickle_dump_to_file, pickle_load_from_file)


def main(argv):
  # program.py <input filepath> <output filepath>

  # Use pickling_support to pickle traceback.
  # https://github.com/ionelmc/python-tblib#pickling-tracebacks
  pickling_support.install()

  try:
    assert len(argv) == 3, len(argv)
    input_filepath = argv[1]
    output_filepath = argv[2]
    func, args, kwargs = pickle_load_from_file(input_filepath)
  except Exception:
    logging.exception('Error on reading pickle')
    return 1

  result = None
  try:
    res = func(*args, **kwargs)
    result = (res, None)
    return 0
  except Exception:
    result = (None, sys.exc_info())
    return 0
  finally:
    pickle_dump_to_file(result, output_filepath)


if __name__ == "__main__":
  # Exit with 0 on target function success and failure.
  #           1 on runner error.
  sys.exit(main(sys.argv))
